import pandas as pd
from pytdx.hq import TdxHq_API

class handlePytdx():

    def get_data(self, stock_code):
        # 创建API对象
        api = TdxHq_API()

        if(int(stock_code) > 60000):
            market = 1
        else:
            market = 0
        
        # K线种类
        # 0 5分钟K线 
        # 1 15分钟K线 
        # 2 30分钟K线 
        # 3 1小时K线 
        # 4 日K线
        # 5 周K线
        # 6 月K线
        # 7 1分钟
        # 8 1分钟K线 
        # 9 日K线
        # 10 季K线
        # 11 年K线
        kline_type = 9

        # 连接到通达信服务器（示例IP和端口，实际使用时需替换为有效的服务器地址和端口）
        # 124.70.133.119  47.116.105.28 121.36.54.217 121.36.81.195 123.249.15.60
        with api.connect('124.70.133.119', 7709):
            # 获取股票日线数据
            # 参数解释：市场代码（0：深圳，1：上海），股票代码，K线种类（9：日线），开始位置，请求数量
            data = api.to_df(api.get_security_bars(kline_type, market, stock_code, 0, 10))
            
            # 将日期列转换为日期格式
            data['Date'] = pd.to_datetime(data['datetime'])

            # 设置日期列为索引
            data.set_index('Date', inplace=True)
            return data